use dashmap::DashMap;
use log::warn;
use tokio::sync::{mpsc, oneshot};

use super::proto::Frame;

pub enum Tx {
    One(oneshot::Sender<Frame>),
    Mpsc(mpsc::Sender<Frame>),
}

pub struct Router {
    m: DashMap<u32, Tx>,
}

impl Router {

    pub fn new() -> Self {
        Self { m: DashMap::new() }
    }

    pub fn notify(&self, frame: Frame) {
        let id = ((frame.src() as u32) << 16) | ((frame.ty() as u32) << 8) | (frame.seq() as u32);
        let tx = self.m.remove(&id);
        match tx {
            Some((_, Tx::One(tx))) => {
                _ = tx.send(frame);
            },
            Some((_, Tx::Mpsc(tx))) => {
                _ = tx.try_send(frame);
            },
            None => {
                warn!("router notify no receiver: {}-{}-{}", frame.src(), frame.ty(), frame.seq());
            }
        };
        
    }

    pub fn subscribe_one(&self, src: u8, ty: u8, seq: u8) -> oneshot::Receiver<Frame> {
        let id = ((src as u32) << 16) | ((ty as u32) << 8) | (seq as u32);
        let (tx, rx) = oneshot::channel();
        self.m.insert(id, Tx::One(tx));
        rx
    }

    pub fn subscribe_mpsc(&self, src: u8, ty: u8, seq: u8) -> mpsc::Receiver<Frame> {
        let id = ((src as u32) << 16) | ((ty as u32) << 8) | (seq as u32);
        let (tx, rx) = mpsc::channel(1);
        self.m.insert(id, Tx::Mpsc(tx));
        rx
    }
}
